"""
=========================================================
Grover 算法框架（Amplitude Amplification Framework for cqlib）
=========================================================

模块功能：
-----------
该模块构建了一个完整的 Grover 幅度放大执行系统，包含：

1️ 问题定义 (`AmplificationProblem`)
    - Oracle（标记“好态”）
    - 状态制备电路 A
    - “好态判定函数” is_good_state()
    - 后处理函数 post_processing()

2️ 算法执行器 (`Grover`)
    - 自动构建 Q = A·S₀·A†·S_f
    - 支持固定/增长式迭代次数
    - 调用 cqlib 的模拟器采样测量
    - 统计每轮概率分布与最优结果

3️ 结果封装 (`GroverResult`)

核心流程：
-----------
|ψ₀⟩  --A-->  |s⟩  --(Q)^k-->  (放大目标态概率)  --测量-->

其中：
    - Q = A·S₀·A†·S_f （Grover 扩散算子）
    - S_f：Oracle，对目标态加 -1 相位
    - S₀：零反射 about |0…0⟩
    - A：状态制备电路（通常 H^{⊗n}）
"""

from __future__ import annotations
from typing import Callable, List, Union, Optional, Any, Iterator
import numpy as np
import random
import itertools

from circuits.algorithm_circuit import Algorithm_Circuit as Circuit
from cqlib.circuits.qubit import Qubit
from cqlib.simulator import StatevectorSimulator

from algorithm.grover_operator import grover_operator


# ============================================================
# 位序重排辅助函数
# ============================================================
def _remap_bits_from_sorted(label_sorted: str, measured_qubits: List[int], target_order: List[int]) -> str:
    """
    把 cqlib 模拟器返回的测量结果位序（升序排列）重新映射回用户定义的目标位序。

    举例：
        measured_qubits = [2, 0]
        label_sorted = '01'（对应 qubit 0=1, qubit 2=0）
        → 返回 '10'（按用户顺序 [2,0] 排）

    参数：
        label_sorted   : 升序测量比特结果（模拟器输出）
        measured_qubits: 实际测量的比特索引
        target_order   : 用户希望呈现的位序（通常与 measured_qubits 相同）

    返回：
        调整后的位串 label_target
    """
    mq_sorted = sorted(measured_qubits)
    bit_at = {mq_sorted[i]: label_sorted[i] for i in range(len(mq_sorted))}
    return ''.join(bit_at[q] for q in target_order)


# ============================================================
# 幅度放大问题定义类
# ============================================================
class AmplificationProblem:
    """
    定义“Grover 幅度放大问题”的通用接口。

    参数：
        oracle:           Circuit，Oracle 电路（标记目标态）
        state_preparation: 状态制备电路 A（若为 None，则默认为 H^{⊗n}）
        objective_qubits: 需要测量的比特索引（None = 全部）
        is_good_state:    判定函数 f(bitstring) → bool
        post_processing:  对测量结果的后处理函数

    在 Grover 算法中，用户只需定义：
        - oracle
        - is_good_state()
    其余部分（A、S₀、Q、模拟）由算法框架自动完成。
    """
    def __init__(
        self,
        oracle: Circuit,
        state_preparation: Optional[Circuit] = None,
        objective_qubits: Union[int, List[int], None] = None,
        is_good_state: Optional[Callable[[str], bool]] = None,
        post_processing: Optional[Callable[[str], Any]] = None
    ):
        self._oracle = oracle
        self._state_preparation = state_preparation
        self._objective_qubits = objective_qubits
        self._is_good_state = is_good_state
        self._post_processing = post_processing

    # ---------- 属性访问 ----------
    @property
    def oracle(self) -> Circuit:
        """返回 Oracle 电路 S_f"""
        return self._oracle

    @property
    def state_preparation(self) -> Circuit:
        """返回状态制备电路 A；若未定义则默认全 H。"""
        if self._state_preparation is None:
            n = self.oracle.num_qubits
            circ = Circuit(n)
            for q in range(n):
                circ.h(q)
            return circ
        return self._state_preparation

    @property
    def objective_qubits(self) -> List[int]:
        """返回需要测量的目标比特列表。"""
        if self._objective_qubits is None:
            return list(range(self.oracle.num_qubits))
        if isinstance(self._objective_qubits, int):
            return [self._objective_qubits]
        return list(self._objective_qubits)

    @property
    def is_good_state(self) -> Callable[[str], bool]:
        """返回目标态判定函数 f(bitstring) → bool。"""
        if self._is_good_state is None:
            raise ValueError("必须提供 is_good_state 函数")
        return self._is_good_state

    @property
    def post_processing(self) -> Callable[[str], Any]:
        """返回后处理函数（若无则恒等）。"""
        return self._post_processing or (lambda x: x)


# ============================================================
# 结果封装类
# ============================================================
class GroverResult:
    """
    存储一次 Grover 放大执行的全部结果。
    """
    def __init__(self):
        self.iterations: Optional[List[int]] = None           # 实际迭代次数列表
        self.top_measurement: Optional[str] = None            # 概率最高的测量结果
        self.assignment: Optional[Any] = None                 # 后处理结果
        self.oracle_evaluation: Optional[bool] = None         # 是否符合好态判定
        self.circuit_results: Optional[List[dict[str, float]]] = None  # 每轮测量概率分布
        self.max_probability: Optional[float] = None          # 最高测量概率


# ============================================================
# Grover 算法执行类
# ============================================================
class Grover:
    """
    Grover 算法执行器，支持多种迭代策略与模拟执行。

    参数：
        iterations:           固定迭代次数或自定义序列（int/list/iterable）
        growth_rate:          若指定，则按 r^k 增长（例如 1.5^k）
        sample_from_iterations: 是否随机采样 [0, t] 中的一轮迭代次数
        shots:                模拟采样次数（每轮）

    限制：
        - iterations 与 growth_rate 不能同时使用
        - growth_rate 必须 > 1
    """
    def __init__(
        self,
        iterations: Union[int, List[int], Iterator[int], None] = None,
        growth_rate: float | None = None,
        sample_from_iterations: bool = False,
        shots: int = 1024
    ):
        # 参数合法性检查
        if growth_rate is not None and iterations is not None:
            raise ValueError("不能同时指定 iterations 和 growth_rate")

        # 动态生成迭代次数序列
        if growth_rate is not None:
            if growth_rate <= 1:
                raise ValueError("growth_rate 必须大于 1")
            self._iterations = (int(growth_rate ** x) for x in itertools.count(1))
        elif isinstance(iterations, int):
            self._iterations = [iterations]
        else:
            self._iterations = iterations if iterations is not None else []

        self._sample_from_iterations = sample_from_iterations
        self._shots = shots

    # ---------- 静态方法：最佳迭代次数估计 ----------
    @staticmethod
    def optimal_num_iterations(num_solutions: int, num_qubits: int) -> int:
        """
        根据理论公式估算最佳迭代次数：
            k ≈ round(arccos(√p) / (2·arcsin(√p)))
        其中 p = num_solutions / 2^n。
        """
        amp = np.sqrt(num_solutions / 2**num_qubits)
        return round(np.arccos(amp) / (2 * np.arcsin(amp)))

    # ---------- 构建单个 Grover 电路 ----------
    def construct_circuit(
        self,
        problem: AmplificationProblem,
        power: int,
        measurement: bool = False
    ) -> Circuit:
        """
        构造一次 Grover 电路：
            |ψ₀⟩ --A--> |s⟩ --(Q)^power--> |ψ_k⟩ --(Measure)
        """
        n = problem.oracle.num_qubits
        qc = Circuit(n)

        # (1) 状态制备 A
        qc = qc.compose(problem.state_preparation)

        # (2) 扩散算子 Q^power
        if power > 0:
            grover_op = grover_operator(problem.oracle, problem.state_preparation)
            for _ in range(power):
                qc = qc.compose(grover_op)

        # (3) 测量指定目标量子位
        if measurement:
            qc.measure(problem.objective_qubits)

        return qc

    # ---------- 主函数：执行放大 ----------
    def amplify(self, problem: AmplificationProblem) -> GroverResult:
        """
        执行 Grover 幅度放大流程并返回结果对象。
        自动控制迭代次数、采样、结果提取。
        """
        res = GroverResult()
        iterations_used: List[int] = []
        all_circuit_results: List[dict[str, float]] = []

        max_prob = 0.0
        top_label_target_order = "0" * len(problem.objective_qubits)
        oracle_ok = False

        # 上限防止无限迭代
        max_iterations = max(10, 2 ** problem.oracle.num_qubits)
        max_power = int(np.ceil(2 ** (len(problem.objective_qubits) / 2)))

        # 获取迭代次数迭代器
        if isinstance(self._iterations, list):
            iterator = iter(self._iterations)
        else:
            iterator = self._iterations

        # ========== 主循环 ==========
        for _ in range(max_iterations):
            try:
                power = next(iterator)
            except StopIteration:
                break
            if power > max_power:
                break

            iterations_used.append(power)
            # 若启用采样模式，则随机选取 [0, power] 中的迭代次数
            t_power = random.randint(0, power) if self._sample_from_iterations else power

            # 构造电路并执行模拟
            qc = self.construct_circuit(problem, t_power, measurement=True)

            # 调用 cqlib 状态向量模拟器
            sim = StatevectorSimulator(qc)
            counts_sorted = sim.sample(shots=self._shots, is_sorted=True)  # {'00': N, '01': M, ...}

            # 位序重排 + 归一化
            measured = problem.objective_qubits
            total = sum(counts_sorted.values())
            counts_target: dict[str, int] = {}
            for label_sorted, c in counts_sorted.items():
                label_target = _remap_bits_from_sorted(label_sorted, measured_qubits=measured, target_order=measured)
                counts_target[label_target] = counts_target.get(label_target, 0) + c
            probs_target = {k: v / total for k, v in counts_target.items()}
            all_circuit_results.append(probs_target)

            # 找出概率最高的结果
            top_label_target_order, max_prob = max(probs_target.items(), key=lambda x: x[1])

            # 若定义了好态判定函数，则检查是否成功
            if problem._is_good_state is not None:
                oracle_ok = problem.is_good_state(top_label_target_order)
                if oracle_ok:
                    break

        # ======= 汇总结果 =======
        res.iterations = iterations_used
        res.top_measurement = top_label_target_order
        res.assignment = problem.post_processing(top_label_target_order)
        res.oracle_evaluation = oracle_ok
        res.circuit_results = all_circuit_results
        res.max_probability = float(max_prob)
        return res
